"""
Copyright (c) 2024 Beijing Jiaotong University
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Yunzhou Tang
Created: 2024/02/14
Supported by: National Key Research and Development Program of China
"""

from . import *
import numpy as np

def transmitter_B2B(
    num_symbols: int = 2**16,
    bits_per_symbol: int = 4,
    total_baud: float = 40e9,
    up_sampling_factor: int = 4,
    roll_off: float = 0.01,
    osnr: int = 32,
    power: int = 0
):

    """ 首先产生发射端X/Y双偏振信号 """
    signal_bits = gen_bits(num_symbols, bits_per_symbol)

    """ 生成双偏振符号序列 """ 
    signals = modulation(signal_bits, bits_per_symbol)

    """ 均衡的训练序列 """
    prev_symbols = signals
    
    """ 上采样 """
    signals = up_sample(signals, up_sampling_factor)

    """ RRC """
    signals = pulse_shaper(signals, up_sampling_factor,
                           roll_off, total_baud)

    """ 根据OSNR添加高斯白噪声 """
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    signals = gaussian_noise(signals, osnr, sampling_rate)

    #处理信号进入光纤计算
    P_dB = power   #dBm
    power = pow(10, (P_dB / 10)) * 1e-3
    TX_X = signals[0].reshape(-1)
    TX_Y = signals[1].reshape(-1)
    TX_X = np.divide(TX_X, np.sqrt(np.mean(np.square(np.abs(TX_X)))))
    TX_Y = np.divide(TX_Y, np.sqrt(np.mean(np.square(np.abs(TX_Y)))))
    TX_X = TX_X * np.sqrt(power)
    TX_Y = TX_Y * np.sqrt(power)
    signals_ = [TX_X, TX_Y]
    signals_ = np.expand_dims(signals_, axis=2)
 
    signals_ = np.array(signals_)
    signals_ = signals_.transpose()
    signals_ = signals_[0]

    #将待传输信号与均衡的训练序列信号合并为一个变量 signals长度为4 前两个位置存放待传输信号 后两个位置存放均衡的训练序列信号
  
    return [signals_, prev_symbols[0], prev_symbols[1]]

